from copy import deepcopy

import matplotlib.pyplot as plt
import numpy as np
import torch
from MinkowskiEngine import SparseTensor
from sklearn.manifold import TSNE
from tqdm import tqdm

from utils.metrics import compute_IoU

CLASSES_NUSCENES = [
    "barrier",
    "bicycle",
    "bus",
    "car",
    "construction_vehicle",
    "motorcycle",
    "pedestrian",
    "traffic_cone",
    "trailer",
    "truck",
    "driveable_surface",
    "other_flat",
    "sidewalk",
    "terrain",
    "manmade",
    "vegetation",
]

CLASSES_KITTI = [
    "car",
    "bicycle",
    "motorcycle",
    "truck",
    "other-vehicle",
    "person",
    "bicyclist",
    "motorcyclist",
    "road",
    "parking",
    "sidewalk",
    "other-ground",
    "building",
    "fence",
    "vegetation",
    "trunk",
    "terrain",
    "pole",
    "traffic-sign",
]

# CLASSES_NUSCENES = [
#     "noise",
#     "animal",
#     "human.pedestrian.adult",
#     "human.pedestrian.child",
#     "human.pedestrian.construction_worker",
#     "human.pedestrian.personal_mobility",
#     "human.pedestrian.police_officer",
#     "human.pedestrian.stroller",
#     "human.pedestrian.wheelchair",
#     "movable_object.barrier",
#     "movable_object.debris",
#     "movable_object.pushable_pullable",
#     "movable_object.trafficcone",
#     "static_object.bicycle_rack",
#     "vehicle.bicycle",
#     "vehicle.bus.bendy",
#     "vehicle.bus.rigid",
#     "vehicle.car",
#     "vehicle.construction",
#     "vehicle.emergency.ambulance",
#     "vehicle.emergency.police",
#     "vehicle.motorcycle",
#     "vehicle.trailer",
#     "vehicle.truck",
#     "flat.driveable_surface",
#     "flat.other",
#     "flat.sidewalk",
#     "flat.terrain",
#     "static.manmade",
#     "static.other",
#     "static.vegetation",
#     "vehicle.ego",
# ]

CLASSES_NUSCENES = [
    "noise",
    "barrier",
    "bicycle",
    "bus",
    "car",
    "construction_vehicle",
    "motorcycle",
    "pedestrian",
    "traffic_cone",
    "trailer",
    "truck",
    "driveable_surface",
    "other_flat",
    "sidewalk",
    "terrain",
    "manmade",
    "vegetation",
]

# CLASSES_NUSCENES = [
#     "noise",
#     "barrier",
#     "bicycle",
#     "bus",
#     "car",
#     "construction_vehicle",
#     "motorcycle",
#     "pedestrian",
#     "traffic_cone",
#     "trailer",
#     "truck",
#     "driveable_surface",
# ]


def evaluate(model_points, model_classifier, model_da, dataloader, config):
    """
    Function to evaluate the performances of a downstream training.
    It prints the per-class IoU, mIoU and fwIoU.
    """
    model_points.eval()
    model_classifier.eval()
    model_da.eval()
    with torch.no_grad():
        # i = 0
        # full_predictions = []
        full_features = []
        ground_truth = []
        for batch in tqdm(dataloader):
            sparse_input = SparseTensor(batch["sinput_F"], batch["sinput_C"], device=0)
            pnt_feats = model_points(sparse_input).F
            # pnt_feats = model_points(sparse_input)
            # pnt_feats = model_da(pnt_feats).F
            # pnt_probs = model_classifier(pnt_feats)
            # # if config["ignore_index"]:
            # #     pnt_feats[:, config["ignore_index"]] = -1e6

            # torch.cuda.empty_cache()
            # preds = pnt_probs.argmax(1).cpu()
            feats = pnt_feats.cpu().to(torch.float16)
            offset = 0
            for j, lb in enumerate(batch["len_batch"]):
                inverse_indexes = batch["inverse_indexes"][j]
                # predictions = preds[inverse_indexes + offset]
                features = feats[inverse_indexes + offset]
                full_features.append(features)

                # remove the ignored index entirely
                # full_predictions.append(predictions)
                ground_truth.append(deepcopy(batch["evaluation_labels"][j]))
                offset += lb
            # i += j

        full_features = torch.cat(full_features).numpy()
        indices = np.random.choice(full_features.shape[0], 3000, replace=False)
        full_features = full_features[indices]
        full_features = (full_features - full_features.min()) / (
            full_features.max() - full_features.min()
        )
        ground_truth = torch.cat(ground_truth).numpy()
        ground_truth = ground_truth[indices]

        tsne = TSNE(n_components=2, init="pca", random_state=0)
        tsne_features = tsne.fit_transform(full_features)

        # _, ax = plt.subplots(1, 1, figsize=(9, 9))
        # ax.axis('off')
        fig = plt.figure(figsize=(9, 9))
        ax = fig.add_axes([0, 0, 1, 1])
        ax.set_axis_off()
        colors = np.array(
            [
                [100, 150, 245],
                [100, 230, 245],
                [30, 60, 150],
                [80, 30, 180],
                [100, 80, 250],
                [255, 30, 30],
                [255, 40, 200],
                [150, 30, 90],
                [255, 0, 255],
                [255, 150, 255],
                [75, 0, 75],
                [175, 0, 75],
                [255, 200, 0],
                [255, 120, 50],
                [0, 175, 0],
                [135, 60, 0],
                [150, 240, 80],
                [255, 240, 150],
                [255, 0, 0],
            ]
        ).astype(np.float_)
        colors /= 255
        colors = tuple(map(tuple, colors))
        ax.set_prop_cycle(color=colors)
        for i in range(1, 17):
            x, y = tsne_features[ground_truth == i].T
            # 画图
            ax.scatter(x, y, label=CLASSES_NUSCENES[i])
        # plt.legend()
        plt.savefig("tmp.svg", bbox_inches="tight", pad_inches=0, transparent=True)
        exit(0)

        m_IoU, fw_IoU, per_class_IoU = compute_IoU(
            torch.cat(full_predictions),
            torch.cat(ground_truth),
            config["classes"],
            ignore_indexes=config["ignore_indexes"],
        )
        classes = [
            CLASSES_NUSCENES[i]
            for i in range(len(CLASSES_NUSCENES))
            if i not in config["ignore_indexes"]
        ]
        print("Per class IoU:")
        if config["dataset"].lower() == "nuscenes":
            print(
                *[
                    f"{a:20} - {b:.3f}"
                    for a, b in zip(classes, (per_class_IoU).numpy())
                ],
                sep="\n",
            )
        elif config["dataset"].lower() == "kitti":
            print(
                *[
                    f"{a:30} - {b:.3f}"
                    for a, b in zip(CLASSES_KITTI, (per_class_IoU).numpy())
                ],
                sep="\n",
            )
        print()
        print(f"mIoU: {m_IoU}")
        print(f"fwIoU: {fw_IoU}")

    return m_IoU
